'''
Created on Nov 12, 2013

@author: mkiyer
'''
from cStringIO import StringIO
import base64

# system-wide imports
import numpy as np
import matplotlib
matplotlib.use('Agg')
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import matplotlib.gridspec as gridspec
from matplotlib.figure import Figure
import matplotlib.pyplot as plt

# flask imports
from flask import render_template

# local imports
from ssea.base import BOOL_DTYPE, Config
from ssea.kernel import RandomState, ssea_kernel

# global figure
global_fig = plt.figure()

class SSEAData:
    def get_details_table(self, sample_metadata):
        # show details of hits
        rows = [['index', 'sample', 'rank', 'raw_weights', 
                 'transformed_weights', 'running_es', 'core_enrichment']]
        for i,ind in enumerate(self.hit_indexes):
            if self.es < 0:
                is_enriched = int(ind >= self.es_rank)
            else:                
                is_enriched = int(ind <= self.es_rank)
            meta = sample_metadata[self.sample_ids[ind]]
            rows.append([i, meta.name, ind+1, self.raw_weights[ind],
                         self.weights_hit[ind], self.es_run[ind], 
                         is_enriched])
        return rows

def ssea_rerun(sample_ids, counts, size_factors, sample_set, seed, config):
    '''    
    sample_ids: list of unique integer ids
    counts: list of floats
    size_factors: list of floats
    sample_set: SampleSet object
    '''
    # convert sample set to membership vector
    membership = np.empty((len(sample_ids),1), dtype=BOOL_DTYPE)
    membership[:,0] = sample_set.get_array(sample_ids)
    # reproduce previous run
    rng = RandomState(seed)
    (ranks, norm_counts, norm_counts_miss, norm_counts_hit, 
     es_vals, es_ranks, es_runs) = \
        ssea_kernel2(counts, size_factors, membership, rng,
                     resample_counts=False,
                     permute_samples=False,
                     add_noise=True,
                     noise_loc=config.noise_loc, 
                     noise_scale=config.noise_scale,
                     method_miss=config.weight_miss,
                     method_hit=config.weight_hit,
                     method_param=config.weight_param)
    # make object for plotting
    m = membership[ranks,0]
    hit_indexes = (m > 0).nonzero()[0]
    d = SSEAData()
    d.es = es_vals[0]
    d.es_run = es_runs[:,0]
    d.es_rank = es_ranks[0]
    d.hit_indexes = hit_indexes
    d.ranks = ranks
    d.sample_ids = sample_ids[ranks]
    d.raw_weights = norm_counts[ranks]
    d.weights_miss = norm_counts_miss[ranks]
    d.weights_hit = norm_counts_hit[ranks]
    return d

def plot_enrichment(result, sseadata, title):
    fig = global_fig
    fig.clf()
    #fig = Figure()
    gs = gridspec.GridSpec(3, 1, height_ratios=[2,1,1])
    # running enrichment score
    ax0 = fig.add_subplot(gs[0])
    x = np.arange(len(sseadata.es_run))
    y = sseadata.es_run
    p1 = ax0.scatter(result.resample_es_ranks, result.resample_es_vals,
                     c='r', s=25.0, alpha=0.3, edgecolors='none')
    p2 = ax0.scatter(result.null_es_ranks, result.null_es_vals,
                     c='b', s=25.0, alpha=0.3, edgecolors='none')
    ax0.plot(x, y, lw=2, color='k', label='Enrichment profile')
    ax0.axhline(y=0, color='gray')
    ax0.axvline(x=sseadata.es_rank, lw=1, linestyle='--', color='black')
    ax0.set_xlim((0, len(sseadata.es_run)))
    #ax0.set_ylim((-1.0, 1.0))
    ax0.grid(True)
    ax0.set_xticklabels([])
    ax0.set_ylabel('Enrichment score (ES)')
    ax0.set_title(title)
    legend = ax0.legend((p1,p2), ('Resampled ES', 'Null ES'), 'upper right',
                        numpoints=1, scatterpoints=1, 
                        prop={'size': 'xx-small'})
    # The frame is matplotlib.patches.Rectangle instance surrounding the legend.
    frame = legend.get_frame()
    frame.set_linewidth(0)
    # membership in sample set
    ax1 = fig.add_subplot(gs[1])
    ax1.vlines(sseadata.hit_indexes, ymin=0, ymax=1, lw=0.25, 
               color='black', label='Hits')
    ax1.set_xlim((0, len(sseadata.es_run)))
    ax1.set_ylim((0, 1))
    ax1.set_xticks([])
    ax1.set_yticks([])
    ax1.set_xticklabels([])
    ax1.set_yticklabels([])
    ax1.set_ylabel('Set')
    # weights
    ax2 = fig.add_subplot(gs[2])
    # TODO: if hit and miss weights differ add a legend
    ax2.plot(sseadata.weights_miss, color='blue')
    ax2.plot(sseadata.weights_hit, color='red')
    #ax2.plot(weights_hit, color='red')
    ax2.set_xlim((0, len(sseadata.es_run)))
    ax2.set_xlabel('Samples')
    ax2.set_ylabel('Weights')
    # draw
    fig.tight_layout()
    return fig

def plot_null_distribution(es, es_null_bins, null_es_hist):
    fig = global_fig
    fig.clf()
    #fig = Figure()
    # get coords of bars    
    left = es_null_bins[:-1]
    height = null_es_hist
    width = [(r-l) for l,r in zip(es_null_bins[:-1],es_null_bins[1:])]
    # make plot
    ax = fig.add_subplot(1,1,1)
    ax.bar(left, height, width)
    ax.axvline(x=es, linestyle='--', color='black')
    ax.set_title('Random ES distribution')
    ax.set_ylabel('P(ES)')
    # calculate percent neg
    percent_neg = 100.0 * sum(null_es_hist[i] for i in xrange(len(null_es_hist))
                              if es_null_bins[i] < 0)
    percent_neg /= float(sum(null_es_hist))
    ax.set_xlabel('ES (Sets with neg scores: %.1f%%)' % 
                  (percent_neg))
    return fig

# def plot():
#     fig = Figure()
#     axis = fig.add_subplot(1, 1, 1)
#     xs = range(100)
#     ys = [random.randint(1, 50) for x in xs]
#     axis.plot(xs, ys)
#     canvas = FigureCanvas(fig)
#     output = StringIO()
#     canvas.print_png(output)
#     response = make_response(output.getvalue())
#     response.mimetype = 'image/png'
#     return response

def figure_to_base64_data(fig):
    output = StringIO()
    canvas = FigureCanvas(fig)
    canvas.print_png(output)
    img_str = base64.b64encode(output.getvalue())
    output.close()
    return img_str
    #fig.savefig(output, format='png')
    #img_str = 'data:image/png;base64,'
    #img_str += base64.b64encode(output.getvalue())
    #html_str = '<img src="%s"/>' % (img_str)
    #return html_str    
    #data = output.getvalue().encode('base64').strip()
    #img_tag = '<img src="data:image/png;base64,{}"/>'
    #return img_tag.format(data) 

#<img src="data:image/png;base64,%s"/>
#...more text and html...
#</body></html>""" % sio.getvalue().encode("base64").strip()

# @app.route('/plot.png')
# def plot():
#     fig = Figure()
#     axis = fig.add_subplot(1, 1, 1)
#     xs = range(100)
#     ys = [random.randint(1, 50) for x in xs]
#     axis.plot(xs, ys)
#     canvas = FigureCanvas(fig)
#     output = StringIO()
#     canvas.print_png(output)
#     response = make_response(output.getvalue())
#     response.mimetype = 'image/png'
#     return response

def create_detailed_report(result, rowmeta, colmeta, sample_set, bcm, config):
    '''
    generate detailed report html
    '''
    # read counts from BigCountMatrix
    counts = np.array(bcm.counts[result.t_id,:], dtype=np.float)
    # remove 'nan' values        
    sample_ids = np.isfinite(counts).nonzero()[0]
    counts = counts[sample_ids]
    size_factors = bcm.size_factors[sample_ids]    
    # rerun ssea
    sseadata = ssea_rerun(sample_ids, counts, size_factors, sample_set, 
                          result.rand_seed, config)
    # enrichment plot
    title = 'Enrichment plot: %s vs. %s' % (rowmeta.name, sample_set.name)
    eplot_figure = plot_enrichment(result, sseadata, title=title)
    eplot_data = figure_to_base64_data(eplot_figure)  
    # null distribution plot
    nplot_figure = plot_null_distribution(result.es, 
                                          Config.NULL_ES_BINS,
                                          result.null_es_hist)
    nplot_data = figure_to_base64_data(nplot_figure)
    # get sample details rows
    details_rows = sseadata.get_details_table(colmeta)
    # render to html
    return render_template('details.html',
                           result=result,
                           sseadata=sseadata,
                           details=details_rows,
                           eplot=eplot_data,
                           nplot=nplot_data,
                           rowmeta=rowmeta,
                           sample_set=sample_set)
